{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Closed-Form Policy to Play LunarLander-v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import logging\n",
    "import itertools\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import gym\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "        format='%(asctime)s [%(levelname)s] %(message)s',\n",
    "        stream=sys.stdout, datefmt='%H:%M:%S')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:00:00 [INFO] env: <LunarLander<LunarLander-v2>>\n",
      "00:00:00 [INFO] action_space: Discrete(4)\n",
      "00:00:00 [INFO] observation_space: Box(-inf, inf, (8,), float32)\n",
      "00:00:00 [INFO] reward_range: (-inf, inf)\n",
      "00:00:00 [INFO] metadata: {'render.modes': ['human', 'rgb_array'], 'video.frames_per_second': 50}\n",
      "00:00:00 [INFO] _max_episode_steps: 1000\n",
      "00:00:00 [INFO] _elapsed_steps: None\n",
      "00:00:00 [INFO] id: LunarLander-v2\n",
      "00:00:00 [INFO] entry_point: gym.envs.box2d:LunarLander\n",
      "00:00:00 [INFO] reward_threshold: 200\n",
      "00:00:00 [INFO] nondeterministic: False\n",
      "00:00:00 [INFO] max_episode_steps: 1000\n",
      "00:00:00 [INFO] _kwargs: {}\n",
      "00:00:00 [INFO] _env_name: LunarLander\n"
     ]
    }
   ],
   "source": [
    "env = gym.make('LunarLander-v2')\n",
    "env.seed(0)\n",
    "for key in vars(env):\n",
    "    logging.info('%s: %s', key, vars(env)[key])\n",
    "for key in vars(env.spec):\n",
    "    logging.info('%s: %s', key, vars(env.spec)[key])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClosedFormAgent:\n",
    "    def __init__(self, _):\n",
    "        pass\n",
    "\n",
    "    def reset(self, mode=None):\n",
    "        pass\n",
    "\n",
    "    def step(self, observation, _reward, _done):\n",
    "        x, y, v_x, v_y, angle, v_angle, contact_left, contact_right = observation\n",
    "\n",
    "        if contact_left or contact_right: # legs have contact\n",
    "            f_y = -10. * v_y - 1.\n",
    "            f_angle = 0.\n",
    "        else:\n",
    "            f_y = 5.5 * np.abs(x) - 10. * y - 10. * v_y - 1.\n",
    "            f_angle = -np.clip(5. * x + 10. * v_x, -4, 4) + 10. * angle + 20. * v_angle\n",
    "\n",
    "        if np.abs(f_angle) <= 1 and f_y <= 0:\n",
    "            action = 0 # do nothing\n",
    "        elif np.abs(f_angle) < f_y:\n",
    "            action = 2 # main engine\n",
    "        elif f_angle < 0.:\n",
    "            action = 1 # left engine\n",
    "        else:\n",
    "            action = 3 # right engine\n",
    "        return action\n",
    "\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "\n",
    "agent = ClosedFormAgent(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:00:00 [INFO] ==== test ====\n",
      "00:00:00 [DEBUG] test episode 0: reward = 251.10, steps = 175\n",
      "00:00:00 [DEBUG] test episode 1: reward = 263.86, steps = 158\n",
      "00:00:00 [DEBUG] test episode 2: reward = 296.23, steps = 228\n",
      "00:00:00 [DEBUG] test episode 3: reward = 290.20, steps = 201\n",
      "00:00:00 [DEBUG] test episode 4: reward = 309.30, steps = 199\n",
      "00:00:00 [DEBUG] test episode 5: reward = 304.04, steps = 171\n",
      "00:00:00 [DEBUG] test episode 6: reward = 260.22, steps = 199\n",
      "00:00:00 [DEBUG] test episode 7: reward = 277.66, steps = 170\n",
      "00:00:00 [DEBUG] test episode 8: reward = 255.32, steps = 181\n",
      "00:00:00 [DEBUG] test episode 9: reward = 261.03, steps = 168\n",
      "00:00:00 [DEBUG] test episode 10: reward = 296.68, steps = 208\n",
      "00:00:00 [DEBUG] test episode 11: reward = 300.02, steps = 222\n",
      "00:00:00 [DEBUG] test episode 12: reward = 303.89, steps = 234\n",
      "00:00:00 [DEBUG] test episode 13: reward = 255.41, steps = 178\n",
      "00:00:00 [DEBUG] test episode 14: reward = 267.43, steps = 180\n",
      "00:00:00 [DEBUG] test episode 15: reward = 309.62, steps = 264\n",
      "00:00:01 [DEBUG] test episode 16: reward = 264.02, steps = 214\n",
      "00:00:01 [DEBUG] test episode 17: reward = 297.51, steps = 216\n",
      "00:00:01 [DEBUG] test episode 18: reward = 261.13, steps = 195\n",
      "00:00:01 [DEBUG] test episode 19: reward = 260.21, steps = 204\n",
      "00:00:01 [DEBUG] test episode 20: reward = 266.38, steps = 248\n",
      "00:00:01 [DEBUG] test episode 21: reward = 253.89, steps = 210\n",
      "00:00:01 [DEBUG] test episode 22: reward = 254.27, steps = 186\n",
      "00:00:01 [DEBUG] test episode 23: reward = 198.83, steps = 604\n",
      "00:00:01 [DEBUG] test episode 24: reward = 289.12, steps = 225\n",
      "00:00:01 [DEBUG] test episode 25: reward = 272.47, steps = 191\n",
      "00:00:01 [DEBUG] test episode 26: reward = 316.36, steps = 267\n",
      "00:00:02 [DEBUG] test episode 27: reward = 235.75, steps = 451\n",
      "00:00:02 [DEBUG] test episode 28: reward = 289.84, steps = 211\n",
      "00:00:02 [DEBUG] test episode 29: reward = 282.75, steps = 194\n",
      "00:00:02 [DEBUG] test episode 30: reward = 249.47, steps = 199\n",
      "00:00:02 [DEBUG] test episode 31: reward = 257.96, steps = 304\n",
      "00:00:02 [DEBUG] test episode 32: reward = 246.41, steps = 243\n",
      "00:00:02 [DEBUG] test episode 33: reward = 285.99, steps = 228\n",
      "00:00:02 [DEBUG] test episode 34: reward = 239.66, steps = 165\n",
      "00:00:02 [DEBUG] test episode 35: reward = 283.29, steps = 210\n",
      "00:00:02 [DEBUG] test episode 36: reward = 244.06, steps = 170\n",
      "00:00:02 [DEBUG] test episode 37: reward = 255.37, steps = 208\n",
      "00:00:02 [DEBUG] test episode 38: reward = 261.29, steps = 171\n",
      "00:00:02 [DEBUG] test episode 39: reward = 253.82, steps = 176\n",
      "00:00:02 [DEBUG] test episode 40: reward = 302.92, steps = 237\n",
      "00:00:02 [DEBUG] test episode 41: reward = 260.70, steps = 167\n",
      "00:00:02 [DEBUG] test episode 42: reward = 255.26, steps = 199\n",
      "00:00:03 [DEBUG] test episode 43: reward = 237.06, steps = 179\n",
      "00:00:03 [DEBUG] test episode 44: reward = 283.01, steps = 148\n",
      "00:00:03 [DEBUG] test episode 45: reward = 258.90, steps = 183\n",
      "00:00:03 [DEBUG] test episode 46: reward = 294.71, steps = 251\n",
      "00:00:03 [DEBUG] test episode 47: reward = 284.71, steps = 161\n",
      "00:00:03 [DEBUG] test episode 48: reward = 281.09, steps = 251\n",
      "00:00:03 [DEBUG] test episode 49: reward = 238.25, steps = 192\n",
      "00:00:03 [DEBUG] test episode 50: reward = 301.94, steps = 240\n",
      "00:00:03 [DEBUG] test episode 51: reward = 263.24, steps = 175\n",
      "00:00:03 [DEBUG] test episode 52: reward = 260.24, steps = 173\n",
      "00:00:03 [DEBUG] test episode 53: reward = 309.02, steps = 196\n",
      "00:00:03 [DEBUG] test episode 54: reward = 294.00, steps = 167\n",
      "00:00:03 [DEBUG] test episode 55: reward = 246.90, steps = 180\n",
      "00:00:03 [DEBUG] test episode 56: reward = 262.93, steps = 215\n",
      "00:00:03 [DEBUG] test episode 57: reward = 285.19, steps = 206\n",
      "00:00:03 [DEBUG] test episode 58: reward = 263.72, steps = 193\n",
      "00:00:03 [DEBUG] test episode 59: reward = -7.60, steps = 102\n",
      "00:00:03 [DEBUG] test episode 60: reward = 260.06, steps = 177\n",
      "00:00:03 [DEBUG] test episode 61: reward = 280.44, steps = 211\n",
      "00:00:03 [DEBUG] test episode 62: reward = 266.07, steps = 175\n",
      "00:00:03 [DEBUG] test episode 63: reward = 272.33, steps = 155\n",
      "00:00:04 [DEBUG] test episode 64: reward = 302.49, steps = 205\n",
      "00:00:04 [DEBUG] test episode 65: reward = 293.76, steps = 227\n",
      "00:00:04 [DEBUG] test episode 66: reward = 280.02, steps = 163\n",
      "00:00:04 [DEBUG] test episode 67: reward = 282.13, steps = 202\n",
      "00:00:04 [DEBUG] test episode 68: reward = -215.60, steps = 308\n",
      "00:00:04 [DEBUG] test episode 69: reward = 306.33, steps = 263\n",
      "00:00:04 [DEBUG] test episode 70: reward = 290.63, steps = 151\n",
      "00:00:04 [DEBUG] test episode 71: reward = 290.89, steps = 177\n",
      "00:00:04 [DEBUG] test episode 72: reward = 271.69, steps = 207\n",
      "00:00:04 [DEBUG] test episode 73: reward = 266.11, steps = 161\n",
      "00:00:04 [DEBUG] test episode 74: reward = 284.49, steps = 146\n",
      "00:00:04 [DEBUG] test episode 75: reward = 266.26, steps = 242\n",
      "00:00:04 [DEBUG] test episode 76: reward = 278.96, steps = 363\n",
      "00:00:05 [DEBUG] test episode 77: reward = 270.56, steps = 231\n",
      "00:00:05 [DEBUG] test episode 78: reward = 269.81, steps = 215\n",
      "00:00:05 [DEBUG] test episode 79: reward = 293.34, steps = 222\n",
      "00:00:05 [DEBUG] test episode 80: reward = 261.76, steps = 191\n",
      "00:00:05 [DEBUG] test episode 81: reward = 277.13, steps = 221\n",
      "00:00:05 [DEBUG] test episode 82: reward = 306.04, steps = 187\n",
      "00:00:05 [DEBUG] test episode 83: reward = 276.49, steps = 219\n",
      "00:00:05 [DEBUG] test episode 84: reward = 284.92, steps = 220\n",
      "00:00:05 [DEBUG] test episode 85: reward = 268.83, steps = 187\n",
      "00:00:05 [DEBUG] test episode 86: reward = 287.16, steps = 178\n",
      "00:00:05 [DEBUG] test episode 87: reward = 279.68, steps = 163\n",
      "00:00:05 [DEBUG] test episode 88: reward = 296.61, steps = 238\n",
      "00:00:05 [DEBUG] test episode 89: reward = 285.85, steps = 154\n",
      "00:00:05 [DEBUG] test episode 90: reward = 280.37, steps = 217\n",
      "00:00:05 [DEBUG] test episode 91: reward = 276.90, steps = 181\n",
      "00:00:05 [DEBUG] test episode 92: reward = 216.17, steps = 291\n",
      "00:00:06 [DEBUG] test episode 93: reward = 19.45, steps = 1000\n",
      "00:00:06 [DEBUG] test episode 94: reward = 259.90, steps = 169\n",
      "00:00:06 [DEBUG] test episode 95: reward = 262.98, steps = 336\n",
      "00:00:06 [DEBUG] test episode 96: reward = 257.21, steps = 259\n",
      "00:00:06 [DEBUG] test episode 97: reward = 239.49, steps = 175\n",
      "00:00:06 [DEBUG] test episode 98: reward = 280.35, steps = 237\n",
      "00:00:06 [DEBUG] test episode 99: reward = 245.76, steps = 194\n",
      "00:00:06 [INFO] average episode reward = 262.72 ± 64.50\n"
     ]
    }
   ],
   "source": [
    "def play_episode(env, agent, max_episode_steps=None, mode=None, render=False):\n",
    "    observation, reward, done = env.reset(), 0., False\n",
    "    agent.reset(mode=mode)\n",
    "    episode_reward, elapsed_steps = 0., 0\n",
    "    while True:\n",
    "        action = agent.step(observation, reward, done)\n",
    "        if render:\n",
    "            env.render()\n",
    "        if done:\n",
    "            break\n",
    "        observation, reward, done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        elapsed_steps += 1\n",
    "        if max_episode_steps and elapsed_steps >= max_episode_steps:\n",
    "            break\n",
    "    agent.close()\n",
    "    return episode_reward, elapsed_steps\n",
    "\n",
    "\n",
    "logging.info('==== test ====')\n",
    "episode_rewards = []\n",
    "for episode in range(100):\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('test episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "logging.info('average episode reward = %.2f ± %.2f',\n",
    "        np.mean(episode_rewards), np.std(episode_rewards))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
